﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Windows.Forms;
using VMExplorer.Win32;

namespace VMExplorer.MemoryObjects
{
  class HeapSegmentObject : RemoteMemoryObject<Native.HEAP_SEGMENT>
  {
    #region Private fields

    HeapObject _parent;

    #endregion
    #region Construction

    public HeapSegmentObject(SafeProcessHandle hProcess, MemoryAddress baseAddress, HeapObject parent)
      : base(hProcess, baseAddress, RegionUsage.Heaps)
    {
      _parent = parent;
    }

    #endregion
    #region Scan rule

    public static void InstallScanRule(Scanner scan)
    {
      scan.NextRegion += delegate(object sender, Scanner.NextRegionEventArgs e)
      {
        if (e.Handled) return;

        Int32 sig = RemoteReader.ReadInt32(e.ProcessHandle, e.Range.Start + 8);
        if (sig == /*0xFFEEFFEE*/-1114130)
        {
          e.MemObjects.Add(new HeapSegmentObject(e.ProcessHandle, e.Range.Start, null));
          e.Handled = true;
        }
      };
    }

    #endregion
    #region IMemoryObject overrides

    public override string Name
    {
      get { return string.Format("{0} - Segment #{1}", _parent.Name, Inner.Entry.SegmentIndex); }
      set { }
    }

    #endregion
  }

  class HeapVirtualAllocdObject : RemoteMemoryObject<Native.HEAP_VIRTUALALLOCATED>
  {
    #region Private fields

    HeapObject _parent;
    int _index;

    #endregion
    #region Construction

    public HeapVirtualAllocdObject(SafeProcessHandle hProcess, MemoryAddress baseAddress, HeapObject parent, int index)
      : base(hProcess, baseAddress, RegionUsage.Heaps)
    {
      _parent = parent;
      _index = index;
    }

    #endregion
    #region IMemoryObject overrides

    public override string Name
    {
      get { return string.Format("{0} - VA #{1}", _parent.Name, _index); }
      set { }
    }

    #endregion
  }

  class HeapObject : RemoteMemoryObject<Native.HEAP>
  {
    #region Private fields

    private List<HeapSegmentObject> _segments = null;
    private List<HeapVirtualAllocdObject> _vallocds = null;
    private List<HeapEntry> _heapEntries = null;

    #endregion
    #region Construction

    private HeapObject(SafeProcessHandle hProcess, MemoryAddress baseAddress, string name)
      : base(hProcess, baseAddress, RegionUsage.Heaps)
    {
      Name = name;
      _segments = new List<HeapSegmentObject>();
      _vallocds = new List<HeapVirtualAllocdObject>();
    }

    public static IEnumerable<IMemoryObject> Create(SafeProcessHandle hProcess, MemoryAddress baseAddress, string name)
    {
      List<IMemoryObject> memObjects = new List<IMemoryObject>();

      if (baseAddress.IsValid)
      {
        #region :: Heap ::

        HeapObject hw = new HeapObject(hProcess, baseAddress, name);
        memObjects.Add(hw);

        #endregion
        #region :: Segments ::

        foreach (MemoryAddress ma in hw.Segments)
        {
          if (ma.IsValid)
          {
            HeapSegmentObject hsw = new HeapSegmentObject(hProcess, ma, hw);
            hw.SegmentList.Add(hsw);
            memObjects.Add(hsw);
          }
        }

        #endregion
        #region :: Virtual allocated blocks ::

        ListEntry virtualAllocdBlocks = new ListEntry(hProcess, hw.Inner.VirtualAllocdBlocks);
        int index = 0;
        foreach (MemoryAddress ma in virtualAllocdBlocks.ForwardList())
        {
          if (ma.IsValid)
          {
            HeapVirtualAllocdObject hvaw = new HeapVirtualAllocdObject(hProcess, ma, hw, index++);
            hw.VallocdList.Add(hvaw);
            memObjects.Add(hvaw);
          }
        }

        #endregion
      }

      return memObjects;
    }

    public static void TryIdentification(SafeProcessHandle hProcess, List<IMemoryObject> memObjects)
    {
      try
      {
        uint exitCode;
        if (RemoteReader.RemoteExecute(hProcess, "kernel32.dll", "GetProcessHeap", out exitCode))
        {
          MemoryAddress heap = new MemoryAddress(new IntPtr((int)exitCode));
          IMemoryObject mo = memObjects
                .OfType<HeapObject>()
                .First(target => target.BaseAddress == heap);

          if (mo != null) mo.Name = "Process Heap";
        }
        if (RemoteReader.RemoteExecute(hProcess, "msvcrt.dll", "_get_heap_handle", out exitCode))
        {
          MemoryAddress heap = new MemoryAddress(new IntPtr((int)exitCode));
          IMemoryObject mo = memObjects
                .OfType<HeapObject>()
                .First(target => target.BaseAddress == heap);

          if (mo != null) mo.Name = "CRT Heap";
        }
        if (RemoteReader.RemoteExecute(hProcess, "mscoree.dll", "GetProcessExecutableHeap", out exitCode))
        {
          MemoryAddress heap = new MemoryAddress(new IntPtr((int)exitCode));
          IMemoryObject mo = memObjects
                .OfType<HeapObject>()
                .First(target => target.BaseAddress == heap);

          if (mo != null) mo.Name = "CLR Native Heap";
        }
      }
      catch
      {
      }

      //  bp kernel32!HeapCreate "gu;.echo;.echo =======================================;?? (void*)@eax;.echo =======================================;kb30;.echo =======================================;.echo.;g"
    }

    #endregion
    #region Scan rule

    public static void InstallScanRule(Scanner scan)
    {
      scan.NextRegion += delegate(object sender, Scanner.NextRegionEventArgs e)
      {
        if (e.Handled) return;

        Int32 sig = RemoteReader.ReadInt32(e.ProcessHandle, e.Range.Start + 8);
        if (sig == /*0xEEFFEEFF*/-285217025)
        {
          e.MemObjects.AddRange(Create(e.ProcessHandle, e.Range.Start, "Unknown heap"));
          e.Handled = true;
        }
      };
    }

    #endregion
    #region IMemoryObject implementation

    public override TabPage[] Tabs
    {
      get
      {
        TabPage tabEntries = new TabPage("Heap entries");
        HeapEntriesTab het = new HeapEntriesTab(this);

        tabEntries.SuspendLayout();
        het.SuspendLayout();

        het.Dock = DockStyle.Fill;

        tabEntries.Controls.Add(het);
        tabEntries.Padding = new System.Windows.Forms.Padding(3);
        tabEntries.UseVisualStyleBackColor = true;

        het.ResumeLayout(false);
        tabEntries.ResumeLayout(false);

        return new TabPage[] { tabEntries };
      }
    }

    #endregion
    #region Public properties

    public List<HeapSegmentObject> SegmentList { get { return _segments; } }

    public List<HeapVirtualAllocdObject> VallocdList { get { return _vallocds; } }

    public List<HeapEntry> Entries
    {
      get
      {
        if (_heapEntries == null)
        {
          _heapEntries = new List<HeapEntry>();

          foreach (HeapSegmentObject hso in SegmentList)
          {
            bool bFirstEntry = true;

            foreach (RegionInfo ri in hso.Ranges)
            {
              MemoryAddress currentChunkBase = bFirstEntry ? hso.BaseAddress : ri.Start;
              HeapEntry hEntry = null;

              if (ri.State == Native.PAGE_STATE.MEM_COMMIT)
              {
                // The very first entry is the segment header
                if (bFirstEntry)
                {
                  hEntry = new HeapEntryRegion(ProcessHandle, currentChunkBase);
                  _heapEntries.Add(hEntry);
                  currentChunkBase += hEntry.ChunkSize;
                  bFirstEntry = false;
                }

                // Iterate over each entries
                do
                {
                  hEntry = HeapEntry.Create(ProcessHandle, currentChunkBase);
                  _heapEntries.Add(hEntry);
                  currentChunkBase += hEntry.ChunkSize;
                } while (!hEntry.IsLastEntry);
              }
              else
              {
                hEntry = new HeapEntryUncommitted(currentChunkBase, (int)ri.Size);
                _heapEntries.Add(hEntry);
              }
            }
          }
        }
        return _heapEntries;
      }
    }

    #endregion

    private MemoryAddress[] Segments
    {
      get
      {
        MemoryAddress[] ma = new MemoryAddress[Inner.LastSegmentIndex + 1];
        for (int i = 0; i <= Inner.LastSegmentIndex; i++)
        {
          ma[i] = new MemoryAddress(Inner.Segments[i]);
        }
        return ma;
      }
    }
  }
}
